﻿// Copyright (c) 2012-2013 fancidev
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

using System;
using System.IO;
using System.IO.Compression;

namespace Util.IO.Compression
{
    /// <summary>
    /// Provides methods to compress and decompress a ZLib stream.
    /// </summary>
    /// <remarks>
    /// This class provides support for reading and writing ZLib streams which
    /// is defined in RFC-1950. A ZLib stream is a wrapper around a Deflate 
    /// stream which is defined in RFC-1951. The ZLib stream adds a two-byte 
    /// header before the Deflate stream and a four-byte checksum after the 
    /// Deflate stream.
    /// 
    /// Since this class uses <code>System.IO.Compression.DeflateStream</code>
    /// to access the compressed stream, there are certain limitations of this
    /// class:
    /// 1 - There is little flexibility in specifying the parameters for the
    ///     compression algorithm, because <code>DeflateStream</code> only
    ///     exposes limited interface to control this algorithm.
    /// 2 - The Adler-32 checksum will not be verified when uncompressing a
    ///     stream, because there is no efficient way to determine on which
    ///     byte <code>DeflateStream</code> finishes using the input (as it
    ///     reads and buffers input in chunks).
    ///     
    /// Reference:
    /// 
    /// MSDN: System.IO.Compression Capabilities.
    /// http://blogs.msdn.com/b/bclteam/archive/2007/05/16/system-io-compression-capabilities-kim-hamilton.aspx
    ///
    /// RFC-1950: ZLIB Compressed Data Format Specification version 3.3.
    /// http://www.ietf.org/rfc/rfc1950.txt
    ///
    /// StackOverflow: Zlib-compatible compression streams.
    /// http://stackoverflow.com/questions/70347/zlib-compatible-compression-streams
    /// </remarks>
    public class ZLibStream : Stream
    {
        private Stream baseStream;
        private CompressionMode mode;
        private Header header;
        private DeflateStream deflateStream;
        private Adler32Checksum checksum;
        private bool leaveOpen;

        /// <summary>
        /// Creates a ZLib stream.
        /// </summary>
        /// <param name="stream">The stream to read or write data.</param>
        /// <param name="mode">Indicates whether to open the stream for 
        /// compression or decompression.</param>
        /// <exception cref="InvalidDataException">The format is not valid.
        /// </exception>
        /// <remarks>
        /// This underlying stream will be automatically disposed when this
        /// object is disposed. However, if the constructor throws an 
        /// exception, the stream will not be automatically disposed.
        /// </remarks>
        public ZLibStream(Stream stream, CompressionMode mode, bool leaveOpen)
        {
            this.baseStream = stream;
            this.mode = mode;
            this.header = (mode == CompressionMode.Compress) ?
                WriteDefaultHeader(stream) : ReadHeader(stream);
            this.deflateStream = new DeflateStream(stream, mode, true);
            this.checksum = new Adler32Checksum();
            this.leaveOpen = leaveOpen;
        }

        public ZLibStream(Stream stream, CompressionMode mode)
            : this(stream, mode, false) { }

        /// <summary>
        /// Reads and validates the two-byte header of a ZLib stream.
        /// </summary>
        private static Header ReadHeader(Stream stream)
        {
            // Read header.
            byte cmf, flg;
            if (!stream.TryReadByte(out cmf))
                throw new InvalidDataException("Missing CMF byte of stream header.");
            if (!stream.TryReadByte(out flg))
                throw new InvalidDataException("Missing FLG byte of stream header.");
            Header header = new Header(cmf, flg);            

            // Validate header.
            if (header.CompressionMethod != CompressionMethod.Deflate)
                throw new InvalidDataException("Compression method is not supported.");
            if (header.HasPresetDictionary)
                throw new InvalidDataException("Preset dictionary is not supported.");
            return header;
        }

        private static Header WriteDefaultHeader(Stream stream)
        {
#if true
            Header header = new Header
            {
                CompressionMethod = CompressionMethod.Deflate,
                HasPresetDictionary = false,
                CompressionLevel = CompressionLevel.Default,
                WindowSize = WindowSize.WS_32K
            };
#else
            Header header = new Header(0x58, 0x85);
#endif
            stream.WriteByte(header.cmf);
            stream.WriteByte(header.flg);
            return header;
        }

        private void CheckDisposed()
        {
            if (deflateStream == null)
                throw new ObjectDisposedException("ZLibStream");
        }

        public override bool CanRead
        {
            get { CheckDisposed(); return deflateStream.CanRead; }
        }

        public override bool CanWrite
        {
            get { CheckDisposed(); return deflateStream.CanWrite; }
        }

        public override bool CanSeek
        {
            get { CheckDisposed(); return deflateStream.CanSeek; }
        }

        public override void Flush()
        {
            CheckDisposed();
            deflateStream.Flush();
        }

        public override long Length
        {
            get { CheckDisposed(); return deflateStream.Length; }
        }

        public override long Position
        {
            get { CheckDisposed(); return deflateStream.Position; }
            set { CheckDisposed(); deflateStream.Position = value; }
        }

        public override int Read(byte[] buffer, int offset, int count)
        {
            CheckDisposed();
            int actual = deflateStream.Read(buffer, offset, count);
            if (actual == 0)
            {
                // EOF. Ideally we should read the 4-byte checksum from the
                // base stream and verify against the checksum we computed.
                // However, DeflateStream reads and buffers the input from
                // the base stream, so the 4 bytes may have already been read
                // into DeflateStream's buffer. Since there is no way to tell
                // where exactly DeflateStream finishes processing the input,
                // we are not able to obtain the checksum unless we receive
                // some extra external parameters such as the stream length.
                // However, that will make the stream interface a bit messy,
                // with little contribution. Therefore, we choose to ignore
                // the checksum when deflate completes.
#if false
                int adler;
                if (!baseStream.TryReadFull(out adler))
                    throw new InvalidDataException("Missing Adler-32 checksum.");
                adler = System.Net.IPAddress.NetworkToHostOrder(adler);
                if (adler != checksum.Value)
                    throw new InvalidDataException("Adler-32 checksum mismatch.");
#endif
                return 0;
            }
            checksum.Update(buffer, offset, actual);
            return actual;
        }

        public override void Write(byte[] buffer, int offset, int count)
        {
            CheckDisposed();
            deflateStream.Write(buffer, offset, count);
            checksum.Update(buffer, offset, count);
        }

        public override long Seek(long offset, SeekOrigin origin)
        {
            CheckDisposed();
            return deflateStream.Seek(offset, origin);
        }

        public override void SetLength(long value)
        {
            CheckDisposed();
            deflateStream.SetLength(value);
        }

        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                if (baseStream != null)
                {
                    deflateStream.Dispose();
                    // Write Adler-32 checksum if in compression mode.
                    if (mode == CompressionMode.Compress)
                    {
                        int adler = System.Net.IPAddress.HostToNetworkOrder(checksum.Value);
                        byte[] buffer = BitConverter.GetBytes(adler);
                        baseStream.Write(buffer, 0, 4);
                    }
                    if (!leaveOpen)
                        baseStream.Dispose();
                }
                this.deflateStream = null;
                this.baseStream = null;
            }
            base.Dispose(disposing);
        }

        /// <summary>
        /// Represents the two-byte header of a ZLib stream.
        /// 
        /// 
        /// The first byte is CMF (Compression Method and flags):
        /// Bits 0 to 3: Compression Method. Always 8 (deflate).
        /// Bits 4 to 7: The base-2 logarithm of the LZ77 window size, minus
        ///              eight. A value of 7 indicates a 32K window size.
        ///              
        /// The second byte is FLG (Flags.
        /// Bits 0 to 4: FCHECK  (check bits for CMF and FLG)
        /// Bit  5     : FDICT   (whether preset dictionary exists)
        /// Bits 6 to 7: FLEVEL  (compression level)
        /// </summary>
        struct Header
        {
            public byte cmf; // first byte
            public byte flg; // second byte

            /// <summary>
            /// Creates a ZLib header with the given CMF and FLG bytes.
            /// </summary>
            public Header(byte cmf, byte flg)
            {
                this.cmf = cmf;
                this.flg = flg;
            }

            /// <summary>
            /// Updates the FCHECK value (lower 5 bites of FLG) so that CMF 
            /// and FLG, when viewed as a 16-bit unsigned integer stored in 
            /// MSB order (CMF*256 + FLG), is a multiple of 31.
            /// </summary>
            private void UpdateFCheck()
            {
                flg &= 0xE0;
                flg |= (byte)(31 - (((int)cmf << 8) | flg) % 31);
            }

            public CompressionMethod CompressionMethod
            {
                get { return (CompressionMethod)(cmf & 0x0F); }
                set
                {
                    cmf = (byte)(cmf & 0xF0 | (byte)value);
                    UpdateFCheck();
                }
            }

            public WindowSize WindowSize
            {
                get { return (WindowSize)(cmf >> 4); }
                set
                {
                    cmf = (byte)((cmf & 0x0F) | ((byte)value << 4));
                    UpdateFCheck();
                }
            }

            public bool HasPresetDictionary
            {
                get { return (flg & 0x20) != 0; }
                set
                {
                    if (value)
                        flg |= 0x20;
                    else
                        flg &= 0xDF;
                    UpdateFCheck();
                }
            }

            public CompressionLevel CompressionLevel
            {
                get { return (CompressionLevel)(flg >> 6); }
                set
                {
                    flg = (byte)((flg & 0x3F) | ((byte)value << 6));
                    UpdateFCheck();
                }
            }
        }

        /// <summary>
        /// Defines the CM (Compression method) field in a ZLib header.
        /// </summary>
        enum CompressionMethod
        {
            Deflate = 8
        }

        /// <summary>
        /// Defines the CINFO (Compression info) field in a ZLib header.
        /// </summary>
        enum WindowSize
        {
            WS_256 = 0,
            WS_512 = 1,
            WS_1K = 2,
            WS_2K = 3,
            WS_4K = 4,
            WS_8K = 5,
            WS_16K = 6,
            WS_32K = 7,
        }

        /// <summary>
        /// Defines the FLEVEL (Compression level) field in a ZLib header.
        /// </summary>
        enum CompressionLevel
        {
            Fastest = 0,
            Fast = 1,
            Default = 2,
            Maximum = 3
        }
    }
}
